/** @addtogroup dft

 *  @{

 */

/*

  Copyright (C) 2016 D Levin (https://www.kfrlib.com)

  This file is part of KFR

  KFR is free software: you can redistribute it and/or modify

  it under the terms of the GNU General Public License as published by

  the Free Software Foundation, either version 3 of the License, or

  (at your option) any later version.

  KFR is distributed in the hope that it will be useful,

  but WITHOUT ANY WARRANTY; without even the implied warranty of

  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the

  GNU General Public License for more details.

  You should have received a copy of the GNU General Public License

  along with KFR.

  If GPL is not suitable for your project, you must purchase a commercial license to use KFR.

  Buying a commercial license is mandatory as soon as you develop commercial activities without

  disclosing the source code of your own applications.

  See https://www.kfrlib.com for details.

 */

#pragma once

#include "../base/complex.hpp"

#include "../base/constants.hpp"

#include "../base/memory.hpp"

#include "../base/read_write.hpp"

#include "../base/small_buffer.hpp"

#include "../base/vec.hpp"

#include "bitrev.hpp"

#include "ft.hpp"

CMT_PRAGMA_GNU(GCC diagnostic push)

#if CMT_HAS_WARNING("-Wshadow")

CMT_PRAGMA_GNU(GCC diagnostic ignored "-Wshadow")

#endif

CMT_PRAGMA_MSVC(warning(push))

CMT_PRAGMA_MSVC(warning(disable : 4100))

namespace kfr

{

template <typename T>

struct dft_stage

{

    size_t stage_size = 0;

    size_t data_size  = 0;

    size_t temp_size  = 0;

    u8* data          = nullptr;

    size_t repeats    = 1;

    size_t out_offset = 0;

    const char* name;

    bool recursion = false;

    void initialize(size_t size) { do_initialize(size); }

    KFR_INTRIN void execute(complex<T>* out, const complex<T>* in, u8* temp) { do_execute(out, in, temp); }

    virtual ~dft_stage() {}

protected:

    virtual void do_initialize(size_t) {}

    virtual void do_execute(complex<T>*, const complex<T>*, u8* temp) = 0;

};

CMT_PRAGMA_GNU(GCC diagnostic push)

#if CMT_HAS_WARNING("-Wassume")

CMT_PRAGMA_GNU(GCC diagnostic ignored "-Wassume")

#endif

namespace internal

{

template <size_t width, bool inverse, typename T>

KFR_SINTRIN cvec<T, width> radix4_apply_twiddle(csize_t<width>, cfalse_t /*split_format*/, cbool_t<inverse>,

                                                const cvec<T, width>& w, const cvec<T, width>& tw)

{

    cvec<T, width> ww  = w;

    cvec<T, width> tw_ = tw;

    cvec<T, width> b1  = ww * dupeven(tw_);

    ww = swap<2>(ww);

    if (inverse)

        tw_ = -(tw_);

    ww      = subadd(b1, ww * dupodd(tw_));

    return ww;

}

template <size_t width, bool use_br2, bool inverse, bool aligned, typename T>

KFR_SINTRIN void radix4_body(size_t N, csize_t<width>, cfalse_t, cfalse_t, cfalse_t, cbool_t<use_br2>,

                             cbool_t<inverse>, cbool_t<aligned>, complex<T>* out, const complex<T>* in,

                             const complex<T>* twiddle)

{

    const size_t N4 = N / 4;

    cvec<T, width> w1, w2, w3;

    cvec<T, width> sum02, sum13, diff02, diff13;

    cvec<T, width> a0, a1, a2, a3;

    a0    = cread<width, aligned>(in + 0);

    a2    = cread<width, aligned>(in + N4 * 2);

    sum02 = a0 + a2;

    a1    = cread<width, aligned>(in + N4);

    a3    = cread<width, aligned>(in + N4 * 3);

    sum13 = a1 + a3;

    cwrite<width, aligned>(out, sum02 + sum13);

    w2 = sum02 - sum13;

    cwrite<width, aligned>(out + N4 * (use_br2 ? 1 : 2),

                           radix4_apply_twiddle(csize_t<width>(), cfalse, cbool_t<inverse>(), w2,

                                                cread<width, true>(twiddle + width)));

    diff02 = a0 - a2;

    diff13 = a1 - a3;

    if (inverse)

    {

        diff13 = (diff13 ^ broadcast<width * 2, T>(T(), -T()));

        diff13 = swap<2>(diff13);

    }

    else

    {

        diff13 = swap<2>(diff13);

        diff13 = (diff13 ^ broadcast<width * 2, T>(T(), -T()));

    }

    w1 = diff02 + diff13;

    cwrite<width, aligned>(out + N4 * (use_br2 ? 2 : 1),

                           radix4_apply_twiddle(csize_t<width>(), cfalse, cbool_t<inverse>(), w1,

                                                cread<width, true>(twiddle + 0)));

    w3 = diff02 - diff13;

    cwrite<width, aligned>(out + N4 * 3, radix4_apply_twiddle(csize_t<width>(), cfalse, cbool_t<inverse>(),

                                                              w3, cread<width, true>(twiddle + width * 2)));

}

template <size_t width, bool inverse, typename T>

KFR_SINTRIN cvec<T, width> radix4_apply_twiddle(csize_t<width>, ctrue_t /*split_format*/, cbool_t<inverse>,

                                                const cvec<T, width>& w, const cvec<T, width>& tw)

{

    vec<T, width> re1, im1, twre, twim;

    split(w, re1, im1);

    split(tw, twre, twim);

    const vec<T, width> b1re = re1 * twre;

    const vec<T, width> b1im = im1 * twre;

    if (inverse)

        return concat(b1re + im1 * twim, b1im - re1 * twim);

    else

        return concat(b1re - im1 * twim, b1im + re1 * twim);

}

template <size_t width, bool splitout, bool splitin, bool use_br2, bool inverse, bool aligned, typename T>

KFR_SINTRIN void radix4_body(size_t N, csize_t<width>, ctrue_t, cbool_t<splitout>, cbool_t<splitin>,

                             cbool_t<use_br2>, cbool_t<inverse>, cbool_t<aligned>, complex<T>* out,

                             const complex<T>* in, const complex<T>* twiddle)

{

    const size_t N4 = N / 4;

    cvec<T, width> w1, w2, w3;

    constexpr bool read_split  = !splitin && splitout;

    constexpr bool write_split = splitin && !splitout;

    vec<T, width> re0, im0, re1, im1, re2, im2, re3, im3;

    split(cread_split<width, aligned, read_split>(in + N4 * 0), re0, im0);